import onnxruntime as ort
import numpy as np

# ----------------------
# 1. 加载 ONNX 模型
# ----------------------
# 可选：指定推理设备（CPU/GPU）
# - CPU: providers=['CPUExecutionProvider']
# - GPU: providers=['CUDAExecutionProvider']（需安装 onnxruntime-gpu）
session = ort.InferenceSession(
    "fsq.onnx",
    providers=['CPUExecutionProvider']  # 根据环境选择
)

# 获取模型输入/输出名称（用于后续传参）
input_name = session.get_inputs()[0].name  # 输入张量的名称
output_name = session.get_outputs()[0].name  # 输出张量的名称


# ----------------------
# 2. 准备输入数据
# ----------------------
# 输入数据需与导出时的 dummy_input 形状一致（如：batch_size=1, seq_len=10, dim=1280）
# 注意：ONNX 通常使用 numpy 数组作为输入，且数据类型需与 PyTorch 一致（如 float32）
input_data = np.random.randn(2, 10, 1280).astype(np.float32)  # 随机生成测试数据


# ----------------------
# 3. 运行推理
# ----------------------
# 传入输入数据（字典格式：{输入名称: 输入数组}）
outputs = session.run([output_name], {input_name: input_data})

# 解析输出（outputs 是一个列表，包含所有输出张量）
quantized_output = outputs[0]  # FSQ 的量化结果（形状与输入对应）
print("输出形状:", quantized_output.shape)  # 例如：(1, 10, 8)（假设降维到 8 维后量化）


